import os
import cv2
import numpy as np
import torch
import torch.nn as nn
from torchvision import models
from DR_project.grad_cam import GradCAM, GradCamPlusPlus
from DR_project.guided_back_propagation import GuidedBackPropagation
from DR_project.utils import circle_crop, get_last_conv_name, norm_image, gen_cam, gen_gb


def interface(img_path, dst_path):
    # 是否使用cuda
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载模型
    model = models.densenet121(pretrained=True)
    fc_feature = model.classifier.in_features
    model.classifier = nn.Linear(fc_feature, 5)
    softmax = nn.Softmax(dim=1)
    model = nn.DataParallel(model).to(device)
    model.load_state_dict(torch.load(
        "/Users/xfdw/Desktop/vueflask/back/DR_project/state-9-0.7958954625621292.pth", map_location='cpu'))

    layer_name = get_last_conv_name(model)  # get last convolution layer
    grad_cam = GradCAM(model, layer_name)  # Grad-CAM
    grad_cam_plus_plus = GradCamPlusPlus(model, layer_name)  # Grad-CAM++

    # 读入并预处理图像
    img_arr = cv2.resize(circle_crop(img_path), (224, 224), interpolation=cv2.INTER_CUBIC) / 255.0
    x = torch.from_numpy(img_arr[:, :, ::-1].astype(np.float32).transpose((2, 0, 1))).unsqueeze(0).to(device)
    x.requires_grad = True
    y = softmax(model(x)).squeeze().cpu().detach().numpy().tolist()  # 模型预测

    # 输出图像
    image_dict = {}
    mask = grad_cam(x, None)  # cam mask
    image_dict['cam'], image_dict['heatmap'] = gen_cam(img_arr[:, :, ::-1].astype(np.float32), mask)
    grad_cam.remove_handlers()
    mask_plus_plus = grad_cam_plus_plus(x, None)  # cam++ mask
    image_dict['cam++'], image_dict['heatmap++'] = gen_cam(img_arr[:, :, ::-1].astype(np.float32), mask_plus_plus)
    grad_cam_plus_plus.remove_handlers()

    # GuidedBackPropagation
    gbp = GuidedBackPropagation(model)
    x.grad.zero_()  # 梯度置零
    grad = gbp(x)

    gb = gen_gb(grad)
    image_dict['gb'] = norm_image(gb)
    # 生成Guided Grad-CAM
    cam_gb = gb * mask[..., np.newaxis]
    image_dict['cam_gb'] = norm_image(cam_gb)

    # print("%%" * 20)
    # print(y, y.index(max(y)))
    # print("%%" * 20)
    # print(image_dict)
    # print("%%" * 20)
    #糖网分级保存图片
    cv2.imwrite(dst_path, image_dict["cam"])

    return y, y.index(max(y)), image_dict
def interface2(img_path, dst_path):
    # 是否使用cuda
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # 加载模型
    model = models.densenet121(pretrained=True)
    fc_feature = model.classifier.in_features
    model.classifier = nn.Linear(fc_feature, 5)
    softmax = nn.Softmax(dim=1)
    model = nn.DataParallel(model).to(device)
    model.load_state_dict(torch.load(
        "/Users/xfdw/Desktop/vueflask/back/DR_project/state-9-0.7958954625621292.pth", map_location='cpu'))

    layer_name = get_last_conv_name(model)  # get last convolution layer
    grad_cam = GradCAM(model, layer_name)  # Grad-CAM
    grad_cam_plus_plus = GradCamPlusPlus(model, layer_name)  # Grad-CAM++

    # 读入并预处理图像
    img_arr = cv2.resize(circle_crop(img_path), (224, 224), interpolation=cv2.INTER_CUBIC) / 255.0
    x = torch.from_numpy(img_arr[:, :, ::-1].astype(np.float32).transpose((2, 0, 1))).unsqueeze(0).to(device)
    x.requires_grad = True
    y = softmax(model(x)).squeeze().cpu().detach().numpy().tolist()  # 模型预测

    # 输出图像
    image_dict = {}
    mask = grad_cam(x, None)  # cam mask
    image_dict['cam'], image_dict['heatmap'] = gen_cam(img_arr[:, :, ::-1].astype(np.float32), mask)
    grad_cam.remove_handlers()
    mask_plus_plus = grad_cam_plus_plus(x, None)  # cam++ mask
    image_dict['cam++'], image_dict['heatmap++'] = gen_cam(img_arr[:, :, ::-1].astype(np.float32), mask_plus_plus)
    grad_cam_plus_plus.remove_handlers()

    # GuidedBackPropagation
    gbp = GuidedBackPropagation(model)
    x.grad.zero_()  # 梯度置零
    grad = gbp(x)

    gb = gen_gb(grad)
    image_dict['gb'] = norm_image(gb)
    # 生成Guided Grad-CAM
    cam_gb = gb * mask[..., np.newaxis]
    image_dict['cam_gb'] = norm_image(cam_gb)

    # print("%%" * 20)
    # print(y, y.index(max(y)))
    # print("%%" * 20)
    # print(image_dict)
    # print("%%" * 20)
    #糖网分级保存图片
    #cv2.imwrite(dst_path, image_dict["cam"])

    return y, y.index(max(y)), image_dict